#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelTriu
{
public:
    __aicore__ inline KernelTriu() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t xShape, uint32_t yShape,
                                int32_t value, uint32_t blockLength, uint32_t tileNum,
                                uint32_t tileLength, uint32_t HignAcc)
    {
        this->value = value;
        this->xShape = xShape;
        this->yShape = yShape;
        this->tileLength = tileLength / BUFFER_NUM;
        this->tileNum = tileNum;
        this->blockLength = blockLength;
        this->HignAcc = HignAcc;
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x, blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)x, blockLength);
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));
    }

    __aicore__ inline void Process()
    {
        if (this->HignAcc){
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int p = 0; p < loopCount; p++)
        {
            CopyIn(p);
            Compute(p);
            CopyOut(p);
        }
        }else{
            lCompute();
        }
    }

private:
    __aicore__ inline void lCompute()
    {
        int32_t dim_xy = this->xShape * this->yShape;
        int32_t dim_f = this->blockLength / dim_xy;
        for (int32_t f = 0; f < dim_f; f++)
        {
            for (int32_t i = 0; i < this->xShape; i++)
            {
                for (int32_t j = 0; j <= i + this->value; j++)
                {
                    yGm.SetValue(f * dim_xy + i * this->yShape + j, 0);
                }
            }
        }
    }

    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[progress * (this->tileLength)], this->tileLength);
        inQueueX.EnQue(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        int32_t start_arr = progress * this->tileLength;
        int32_t xdim = start_arr / this->yShape;
        int32_t rpt = (xdim + this->value) / 128;
        int32_t arr = (xdim + this->value) % 128;
        int32_t arr1 = (arr > 64) ? 64 : arr;
        int32_t arr2 = (arr > 64)?(arr - 64):0;
        Muls(xLocal, xLocal, (DTYPE_X)0, 128 * rpt);
        uint64_t mask[2] = {UINT64_MAX >> arr1, UINT64_MAX > arr2};
        Muls(yLocal, xLocal, (DTYPE_X)0, mask, 1 + rpt, {1, 1, 8, 8});
        outQueueY.EnQue<DTYPE_Y>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * (this->tileLength)], yLocal,
                 (this->tileLength));
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t blockLength;
    int32_t value;
    uint32_t xShape;
    uint32_t yShape;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t HignAcc;
};

extern "C" __global__ __aicore__ void triu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tilingData, tiling);
    // TODO: user kernel impl
    KernelTriu op;
    op.Init(x, y, tilingData.xShape, tilingData.yShape,
            tilingData.value, tilingData.blockLength,
            tilingData.tileNum, tilingData.tileLength,
            tilingData.HignAcc);
    op.Process();
}

